import multiprocessing.queues
import os

# limit the number of cpus used by high performance libraries
import time
from multiprocessing import Value
from multiprocessing.dummy.connection import Connection

import easydict

from cjc_detect.detect_age import create_age_detector
from cjc_detect.detect_sex import create_sex_detector
from cjc_tools import sort_yolov5_opts
from cjc_tools.detect_result_to_label import sex_cls_to_txt, age_cls_to_txt
from cjc_tools.person_counter import PersonCounter
from cjc_tools.pipe_data_structure import create_tracker_2_detector_data
from cjc_tools.save_video import Cv2Mp4Saver, Cv2StreamSaver
from cjc_trackers import traj_map
from cjc_trackers.tracker_tool import put_tracker_data_2, put_tracker_img_2
from project_setting import media_ext

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"

import sys
import numpy as np
from pathlib import Path
import torch
import torch.backends.cudnn as cudnn

FILE = Path(__file__).resolve()
ROOT = FILE.parents[1]  # yolov5 strongsort root directory
WEIGHTS = ROOT / 'weights'
sys.path.append(ROOT)
sys.path.append(WEIGHTS)
sys.path.append(f'{ROOT}/strong_sort/deep/reid')

if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
if str(ROOT / 'yolov5') not in sys.path:
    sys.path.append(str(ROOT / 'yolov5'))  # add yolov5 ROOT to PATH
if str(ROOT / 'strong_sort') not in sys.path:
    sys.path.append(str(ROOT / 'strong_sort'))  # add strong_sort ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

import logging
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.dataloaders import VID_FORMATS, LoadImages, LoadStreams
from yolov5.utils.general import (LOGGER, check_img_size, non_max_suppression, scale_coords, check_requirements, cv2,
                                  check_imshow, xyxy2xywh, increment_path, strip_optimizer, colorstr, print_args,
                                  check_file)
from yolov5.utils.torch_utils import select_device, time_sync
from yolov5.utils.plots import Annotator, colors, save_one_box
from strong_sort.utils.parser import get_config
from strong_sort.strong_sort import StrongSORT

# remove duplicated stream handler to avoid duplicated logging
logging.getLogger().removeHandler(logging.getLogger().handlers[0])


class Tracker:

    def __init__(self, _opt, can_track=1):
        if can_track==1:
            self.can_track = Value("i",1)  # 是否追踪 ，0 False 1
        else:
            self.can_track = can_track
        self.tracker_queue: multiprocessing.queues.Queue = None
        self.vid_cap = None
        self.opt = _opt
        self.person_counter = PersonCounter(0)  # 人计数器
        # self.person_counter.register_new_video_callback(self.newvideo)
        self.conns = []
        self.qs = []
        fp = f'{_opt.project}/webcam'+media_ext
        self.video_saver = Cv2StreamSaver(self.vid_cap, fp)
        self.video_saver.refresh_path()
        self.when_first_person_enter = 0  # 第一个人进入视频
        self.new_record_delay = 3  # 十秒后重新计时
        self.age_detector = create_age_detector()  # 年龄检测器
        self.sex_detector = create_sex_detector()  # 性别检测器
        self.on_video_finished = None  # 视频录制完成后执行
        self.last_video_save_path = self.video_saver.save_path
        self.traj_map = traj_map.Trajectories()
        self.has_data_to_be_saved = False  # 是否已经保存过数据
        self.need_to_refresh_path = False  # 是否已经保存过数据


    def set_video_finished_cb(self, cb):
        self.on_video_finished = cb

    @staticmethod
    def _parse_detect_result(r: list) -> int:
        """得到分类 cls"""
        if r is None:
            return None
        lines = [_l.split(' ')[0] for _l in r]
        return int(lines[0])

    def add_detect_result(self, _id, sex_result, age_result):
        '''
        同时记录性别和年龄
        Parameters
        ----------
        _id
        sex_result
        age_result

        Returns
        -------

        '''
        if _id is None:
            print("id is None", _id)
            return
        sex = self._parse_detect_result(sex_result)
        age = self._parse_detect_result(age_result)
        if sex is not None:
            self.person_counter.sex_group.add_person(int(_id), int(sex))
        if age is not None:
            self.person_counter.age_group.add_person(int(_id), int(age))

    def update_first_person_enter(self):
        self.when_first_person_enter = time.time()

    def how_long_first_person_entered(self):
        return time.time() - self.when_first_person_enter

    def set_conn(self, queue):
        # self.conns.extend(ps)
        self.tracker_queue = queue

    def set_queue(self, q):
        self.qs.extend(q)

    def run_wrapper(self):
        self.run(**self.opt)

    @torch.no_grad()
    def run(self,
            source='0',
            yolo_weights=WEIGHTS / 'yolov5m.pt',  # model.pt path(s),
            strong_sort_weights=WEIGHTS / 'osnet_x0_25_msmt17.pt',  # model.pt path,
            config_strongsort=ROOT / 'strong_sort/configs/strong_sort.yaml',
            imgsz=(640, 640),  # inference size (height, width)
            conf_thres=0.25,  # confidence threshold
            iou_thres=0.45,  # NMS IOU threshold
            max_det=1000,  # maximum detections per image
            device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
            show_vid=False,  # show results
            save_txt=False,  # save results to *.txt
            save_conf=False,  # save confidences in --save-txt labels
            save_crop=False,  # save cropped prediction boxes
            save_vid=False,  # save confidences in --save-txt labels
            nosave=False,  # do not save images/videos
            classes=None,  # filter by class: --class 0, or --class 0 2 3
            agnostic_nms=False,  # class-agnostic NMS
            augment=False,  # augmented inference
            visualize=False,  # visualize features
            update=False,  # update all models
            project=ROOT / 'runs/track',  # save results to project/name
            name='exp',  # save results to project/name
            exist_ok=False,  # existing project/name ok, do not increment
            line_thickness=3,  # bounding box thickness (pixels)
            hide_labels=False,  # hide labels
            hide_conf=False,  # hide confidences
            hide_class=False,  # hide IDs
            half=False,  # use FP16 half-precision inference
            dnn=False,  # use OpenCV DNN for ONNX inference
            ):

        source = str(source)
        save_img = not nosave and not source.endswith('.txt')  # save inference images
        is_file = Path(source).suffix[1:] in (VID_FORMATS)
        is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
        webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
        if is_url and is_file:
            source = check_file(source)  # download

        # Directories
        if not isinstance(yolo_weights, list):  # single yolo model
            exp_name = yolo_weights.stem
        elif type(yolo_weights) is list and len(yolo_weights) == 1:  # single models after --yolo_weights
            exp_name = Path(yolo_weights[0]).stem
        else:  # multiple models after --yolo_weights
            exp_name = 'ensemble'
        exp_name = name if name else exp_name + "_" + strong_sort_weights.stem
        save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok)  # increment run
        (save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

        # Load model
        device = select_device(device)
        model = DetectMultiBackend(yolo_weights, device=device, dnn=dnn, data=None, fp16=half)
        stride, names, pt = model.stride, model.names, model.pt
        imgsz = check_img_size(imgsz, s=stride)  # check image size
        # self.traj_map.reset_canvas(*imgsz)   # 初始化尺寸
        # Dataloader
        if webcam:
            # show_vid = check_imshow()
            cudnn.benchmark = True  # set True to speed up constant image size inference
            dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt)
            nr_sources = len(dataset)
        else:
            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
            nr_sources = 1
        vid_path, vid_writer, txt_path = [None] * nr_sources, [None] * nr_sources, [None] * nr_sources

        # initialize StrongSORT
        cfg = get_config()
        cfg.merge_from_file(config_strongsort)

        # Create as many strong strong_sort instances as there are video sources
        strongsort_list = []
        for i in range(nr_sources):
            strongsort_list.append(
                StrongSORT(
                    strong_sort_weights,
                    device,
                    half,
                    max_dist=cfg.STRONGSORT.MAX_DIST,
                    max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE,
                    max_age=cfg.STRONGSORT.MAX_AGE,
                    n_init=cfg.STRONGSORT.N_INIT,
                    nn_budget=cfg.STRONGSORT.NN_BUDGET,
                    mc_lambda=cfg.STRONGSORT.MC_LAMBDA,
                    ema_alpha=cfg.STRONGSORT.EMA_ALPHA,
                )
            )
            strongsort_list[i].model.warmup()
        outputs = [None] * nr_sources

        # Run tracking
        model.warmup(imgsz=(1 if pt else nr_sources, 3, *imgsz))  # warmup
        dt, seen = [0.0, 0.0, 0.0, 0.0], 0
        curr_frames, prev_frames = [None] * nr_sources, [None] * nr_sources
        print("采集服务器成功启动")
        for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):

            t1 = time_sync()
            im = torch.from_numpy(im).to(device)
            im = im.half() if half else im.float()  # uint8 to fp16/32
            im /= 255.0  # 0 - 255 to 0.0 - 1.0
            if len(im.shape) == 3:
                im = im[None]  # expand for batch dim
            t2 = time_sync()
            dt[0] += t2 - t1

            # Inference
            visualize = increment_path(save_dir / Path(path[0]).stem, mkdir=True) if visualize else False
            pred = model(im, augment=augment, visualize=visualize)
            t3 = time_sync()
            dt[1] += t3 - t2

            # Apply NMS
            pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
            dt[2] += time_sync() - t3

            # Process detections
            for i, det in enumerate(pred):  # detections per image
                if self.can_track.value == 0:
                    continue
                seen += 1
                if webcam:  # nr_sources >= 1
                    p, im0, _ = path[i], im0s[i].copy(), dataset.count
                    p = Path(p)  # to Path
                    s += f'{i}: '
                    txt_file_name = p.name
                    save_path = str(save_dir / p.name)  # im.jpg, vid.mp4, ...
                else:
                    p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
                    p = Path(p)  # to Path
                    # video file
                    if source.endswith(VID_FORMATS):
                        txt_file_name = p.stem
                        save_path = str(save_dir / p.name)  # im.jpg, vid.mp4, ...
                    # folder with imgs
                    else:
                        txt_file_name = p.parent.name  # get folder name containing current img
                        save_path = str(save_dir / p.parent.name)  # im.jpg, vid.mp4, ...
                curr_frames[i] = im0
                self.traj_map.create_canvas_from_bk(im0)  # 设置 traj_map的背景图
                txt_path = str(save_dir / 'tracks' / txt_file_name)  # im.txt
                s += '%gx%g ' % im.shape[2:]  # print string
                imc = im0.copy() if save_crop else im0  # for save_crop

                annotator = Annotator(im0, line_width=2, pil=not ascii)
                if cfg.STRONGSORT.ECC:  # camera motion compensation
                    strongsort_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])

                if det is not None and len(det):
                    self.has_data_to_be_saved = False
                    # Rescale boxes from img_size to im0 size
                    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()

                    # Print results
                    for c in det[:, -1].unique():
                        n = (det[:, -1] == c).sum()  # detections per class
                        s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                    xywhs = xyxy2xywh(det[:, 0:4])
                    confs = det[:, 4]
                    clss = det[:, 5]

                    # pass detections to strongsort
                    t4 = time_sync()
                    outputs[i] = strongsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
                    t5 = time_sync()
                    dt[3] += t5 - t4
                    if self.can_track.value == 0:
                        continue  # 如果 can_track 不可用 则跳过
                    # draw boxes for visualization
                    if len(outputs[i]) > 0:
                        self.update_first_person_enter()
                        self.need_to_refresh_path = True  # 可以刷新了，但是现在还没到时候
                        for j, (output, conf) in enumerate(zip(outputs[i], confs)):

                            bboxes = output[0:4]  # 这个是检测结果，包括位置信息
                            id = output[4]  # 跟踪结果id
                            cls = output[5]  # 分类
                            self.traj_map.draw_point_with_index(bboxes, id)  # 根据id设置颜色
                            if self.can_track.value == 1:
                                self.person_counter.run(id, bboxes)
                            if save_txt:
                                # to MOT format
                                bbox_left = output[0]
                                bbox_top = output[1]
                                bbox_w = output[2] - output[0]
                                bbox_h = output[3] - output[1]
                                # Write MOT compliant results to file
                                with open(txt_path + '.txt', 'a') as f:
                                    f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left,  # MOT format
                                                                   bbox_top, bbox_w, bbox_h, -1, -1, -1, i))

                            if save_vid or save_crop or show_vid:  # Add bbox to image
                                c = int(cls)  # integer class
                                id = int(id)  # integer id
                                label = None if hide_labels else (f'{id} {names[c]}' if hide_conf else \
                                                                      (
                                                                          f'{id} {conf:.2f}' if hide_class else f'{id} {names[c]} {conf:.2f}'))
                                # 画label的行为移动到了if后面

                                if save_crop:
                                    txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''
                                    crop = save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[
                                        c] / f'{id}' / f'{p.stem}.jpg', BGR=True, save=False)
                                    # 检测性别和年龄
                                    sex_result = self.sex_detector.detect(crop, id)
                                    age_result = self.age_detector.detect(crop, id)
                                    sex_txt = sex_cls_to_txt(self._parse_detect_result(sex_result))
                                    age_txt = age_cls_to_txt(self._parse_detect_result(age_result))
                                    label = f'{label} {sex_txt} {age_txt}'
                                    if self.can_track.value == 1:
                                        self.add_detect_result(id, sex_result, age_result)
                                    # if len(self.conns) > 0: 不在使用 connection
                                    #     for c in self.conns:
                                    #         try:
                                    #             c.send(create_tracker_2_detector_data(id, crop))
                                    #         except Exception as e:
                                    #             print(e)
                                    #             # cv2.imshow("dd",crop)
                                    # cv2.waitKey(10)
                                # 这里发送给connection
                                annotator.box_label(bboxes, label, color=colors(c, True))

                    LOGGER.info(f'{s}Done. YOLO:({t3 - t2:.3f}s), StrongSORT:({t5 - t4:.3f}s)')

                else:  # 扫描不到人才进入这个分支
                    strongsort_list[i].increment_ages()
                    if self.has_data_to_be_saved is False and self.need_to_refresh_path and self.how_long_first_person_entered() > self.new_record_delay:
                        self.has_data_to_be_saved = True

                        # if self.on_video_finished is not None : #这段没用了
                        #     if self.last_video_save_path == self.video_saver.save_path:continue
                        #     self.on_video_finished(self.video_saver.save_path,self.person_counter.get_counter_status())
                        #     self.last_video_save_path = self.video_saver.save_path

                    LOGGER.info('No detections')

                # Stream results
                im0 = annotator.result()
                # im0 = self.person_counter.draw_mask(im0)
                # put_tracker_img_2(self.tracker_queue,im0)
                if self.can_track.value == 1:
                    # 如果当前人数是0 则
                    if self.has_data_to_be_saved:  # 正常
                        self.traj_map.set_background(im0)
                        self.video_saver.refresh_path()  # 刷新路径，结束当前录制  一定要先刷新在做其他的。
                        # im0=self.traj_map.result()
                        # put_tracker_data_2(self.tracker_queue,img ,
                        #                    self.person_counter.get_counter_status(),
                        #                    self.last_video_save_path,True)
                        img = self.traj_map.result_to_tracker()
                        self.need_to_refresh_path = False  # 刚刷新完不用再进来了
                        self.has_data_to_be_saved = False  # 该保存的保存完了
                        self.update_first_person_enter()  # 3秒之内不在回来
                        self.last_video_save_path = self.video_saver.save_path
                        put_tracker_data_2(self.tracker_queue, img,
                                           self.person_counter.get_counter_status_for_webapi(),
                                           self.last_video_save_path, True)
                        self.person_counter.clear()  # 前面已经把结果传入了ws 现在清空缓存
                        self.traj_map.clear()  # 清楚轨迹图
                    else:
                        # pass
                        put_tracker_data_2(self.tracker_queue, im0, self.person_counter.get_counter_status(),
                                           self.last_video_save_path)

                if show_vid:
                    cv2.imshow(str(p), im0)
                    cv2.waitKey(1)  # 1 millisecond

                # Save results (image with detections)
                if save_vid:
                    if vid_path[i] != save_path:  # new video
                        vid_path[i] = save_path
                        if isinstance(vid_writer[i], cv2.VideoWriter):
                            vid_writer[i].release()  # release previous video writer
                        if vid_cap:  # video
                            # print("这么获取的", )
                            fps = vid_cap.get(cv2.CAP_PROP_FPS)
                            w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
                            h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
                        else:  # stream
                            # print("这么获取的2", )
                            fps, w, h = 30, im0.shape[1], im0.shape[0]
                        save_path = str(Path(save_path).with_suffix('.mp4'))  # force *.mp4 suffix on results videos
                        vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
                    vid_writer[i].write(im0)
                    if len(self.person_counter.pos_dic) > 0:
                        self.video_saver.write(im0)
                prev_frames[i] = curr_frames[i]

        # Print results
        t = tuple(x / seen * 1E3 for x in dt)  # speeds per image
        LOGGER.info(
            f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms strong sort update per image at shape {(1, 3, *imgsz)}' % t)
        if save_txt or save_vid:
            s = f"\n{len(list(save_dir.glob('tracks/*.txt')))} tracks saved to {save_dir / 'tracks'}" if save_txt else ''
            LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")
        if update:
            strip_optimizer(yolo_weights)  # update model (to fix SourceChangeWarning)
